# -*- coding: utf-8 -*-
import numpy as np
import cv2
import glob
import os





global SAVE_PATH# , mtx_left, dist_left, mtx_right, dist_right, shape, R, T, R1, R2,  P1, P2, Q

SAVE_PATH = 'CorrectFiles'

if not os.path.exists(SAVE_PATH):
    os.mkdir(SAVE_PATH)

def calibrate(path, boardSize, squareSize, left=True):
    print('Start Calibrate')
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    Nx_cor = boardSize[0]
    Ny_cor = boardSize[1]

    objp = np.zeros((Nx_cor * Ny_cor, 3), np.float32)
    objp[:, :2] = np.mgrid[0:Nx_cor * squareSize:squareSize, 0:Ny_cor * squareSize:squareSize].T.reshape(-1, 2)
    objpoints = []  # 3d points in real world space
    imgpoints = []  # 2d points in image plane.
    images = glob.glob(os.path.join(path, '*.jpg'))
    print(images)



    for fname in images:
        frame = cv2.imread(fname)
        # Our operations on the frame come here
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    
        ret, corners = cv2.findChessboardCorners(gray, (Nx_cor, Ny_cor), None)  # Find the corners
        print(ret)
        # If found, add object points, image points
        if ret == True:
            corners = cv2.cornerSubPix(gray, corners, (5, 5), (-1, -1), criteria)
            objpoints.append(objp)
            imgpoints.append(corners)
            cv2.drawChessboardCorners(frame, (Nx_cor, Ny_cor), corners, ret)
            

       
    global mtx, dist

    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
    print(mtx, dist)

    mean_error = 0
    for i in range(len(objpoints)):
        imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
        error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        mean_error += error

    print ("total error: ", mean_error / len(objpoints))
        # # When everything done, release the capture
    if left:
        np.savez('calibrate_left.npz', mtx=mtx, dist=dist)
    else:
        np.savez('calibrate_right.npz', mtx=mtx, dist=dist)
    return mtx, dist

def AutoGetImagePointsInStero(steroPath, boardSize, squareSize):
    '''
        steroPath: 存放图片的根目录
    '''
    print('Start AutogetImagePointsInStero')
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    Nx_cor = boardSize[0]
    Ny_cor = boardSize[1]

    objp = np.zeros((Nx_cor * Ny_cor, 3), np.float32)
    objp[:, :2] = np.mgrid[0:Nx_cor * squareSize:squareSize, 0:Ny_cor * squareSize:squareSize].T.reshape(-1, 2)# objp是三维的
    objpoints = []  # 3d points in real world space
    imgpoints = [[], []]  # 2d points in image plane.
    print(os.path.join(steroPath, 'left'))
    left_images = glob.glob(r'{}/*.jpg'.format(os.path.join(steroPath, 'left')))
    print(left_images)
    imageShape = None
    for fname_left in left_images:
        print(fname_left)
        frame_left = cv2.imread(fname_left)
        fname_right = os.path.join(steroPath, os.path.join('right', os.path.basename(fname_left)))
        if not os.path.exists(fname_right):
            print("{} not in right".format(os.path.basename(fname_right)))
            continue
        frame_right = cv2.imread(fname_right)
        # Our operations on the frame come here
        if not frame_left.shape == frame_right.shape:
            print('{} has different shape in left and right'.format(os.path.basename(fname_right)))
            continue
        gray_left = cv2.cvtColor(frame_left, cv2.COLOR_BGR2GRAY)
        gray_right = cv2.cvtColor(frame_right, cv2.COLOR_BGR2GRAY)
        if not imageShape:
            imageShape = gray_left.shape
        elif gray_left.shape != imageShape:
            print('{} has different shape with before'.format(os.path.basename(fname_right)))
            continue
        print('Start Find Chess left\n')
        ret_left, corners_left = cv2.findChessboardCorners(gray_left, (Nx_cor, Ny_cor), None)  # Find the corners
        print('Start Find Chess right\n')
        ret_right, corners_right = cv2.findChessboardCorners(gray_right, (Nx_cor, Ny_cor), None)
        
        # If found, add object points, image points
        print('End Find Chess\n')
        if ret_left and ret_right:
            print('Start compute super pix\n')
            corners_left = cv2.cornerSubPix(gray_left, corners_left, (5, 5), (-1, -1), criteria)
            corners_right = cv2.cornerSubPix(gray_right, corners_right, (5, 5), (-1, -1), criteria)
            objpoints.append(objp)
            imgpoints[0].append(corners_left)
            imgpoints[1].append(corners_right)
            # show corners
            # cv2.drawChessboardCorners(frame_left, (Nx_cor, Ny_cor), corners_left, ret_left)
            # cv2.imshow('frame', frame_left)
            # k = cv2.waitKey(0)
            # if k == 27:
            #     cv2.destroyAllWindows()
        else:
            print('{} noy find chess'.format(fname_left))
            continue
        print('Finish {}\n'.format(fname_left))
    return objpoints, imgpoints, imageShape


def WriteImagePointsInStero():
    # 现实中标定的坐标点，注意需要以一个点为x,y,z的坐标原点
    gray= cv2.imread(r'20200103Pattern2/left.jpg', 0)
    shape= gray.shape
    objpoints = []

    points = np.array([[0, 0, 0],
                   [0,0,1.4],
                   [0,0,7],
                   [0,0,9.5],
                   [20,10,0],
                   [20,10,1.4],
                   [20,10,7],
                   [20,10,9.5],
                   [10, 30, 0],
                   [10, 30, 7.4],
                   [10, 30, 14.5],
                   [20, 25, 0],
                   [20,25,5.5],
                   [30,15,0],
                   [30,15, 4.1],
                   [30,15,8.3]], dtype=np.float32)
    objpoints.append(points)
    imgpoints = [[], []]

    # 现实的点在图片上的对应
    left_points = np.array([[3781, 1882],
                            [3781, 1788],
                            [3781, 1433],
                            [3781, 1270],
                            [2867, 1850],
                            [2867, 1773],
                            [2867, 1492],
                            [2867, 1368],
                            [1897, 1885],
                            [1897, 1406],
                            [1897, 990],
                            [2135, 1853],
                            [2135, 1566],
                            [2450, 1823],
                            [2450, 1646],
                            [2450, 1460]], dtype=np.float32)
    right_points = np.array([[1893, 1771],
                         [1893, 1596],
                         [1893, 1167],
                         [1893, 981],
                         [1489, 1637],
                         [1489, 1566],
                         [1489, 1267],
                         [1489, 1134],
                         [241, 1658],
                         [241, 1211],
                         [241, 833],
                         [748, 1643],
                         [748, 1362],
                         [1330, 1616],
                         [1330, 1436],
                         [1330, 1241]], dtype=np.float32)
    imgpoints[0].append(left_points)
    imgpoints[1].append(right_points)
    
    return objpoints, imgpoints, shape
    
    
    

def strtoCalibrate(objpoints, imgpoints, mtx_left, dist_left, mtx_right, dist_right, shape):

    stereocalibration_criteria = (cv2.TERM_CRITERIA_MAX_ITER + cv2.TERM_CRITERIA_EPS, 500, 1e-5)
    stereocalibration_flags = cv2.CALIB_FIX_INTRINSIC
    stereocalibration_retval, mtx_left, dist_left, mtx_right, dist_right, R, T, E, F = cv2.stereoCalibrate(objpoints, imgpoints[0], imgpoints[1], mtx_left, dist_left, mtx_right, dist_right, shape[::-1], criteria=stereocalibration_criteria, flags=stereocalibration_flags)
    print('R is {}'.format(R))
    print('T is {}'.format(T))

    if stereocalibration_retval:
        R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(mtx_left, dist_left, mtx_right, dist_right, shape, R, T)
        np.savez('calibrate_stero.npz', R=R, T=T, R1=R1, R2=R2, P1=P1, P2=P2, Q=Q,
                                  shape=shape)
    
    print('P1 is {}'.format(P1))
    print('P2 is {}'.format(P2))
    print('Q is {}'.format(Q))
    # cv2.triangulatePoints()
    return R1, R2, P1, P2, R, T, Q

def correctImage(imageL, imageR, shape, mtx_left, dist_left, R1, mtx_right, dist_right, R2):
    print(imageL)
    print(imageR)
    imgL = cv2.imread(imageL)
    imgR = cv2.imread(imageR)
    h, w = shape
    newcameramtx_left, roi_left=cv2.getOptimalNewCameraMatrix(mtx_left,dist_left,(w,h),1,(w,h))
    # dst = cv2.undistort(imgL, mtx_left, dist_left, R1, newcameramtx_left)
    mapx,mapy = cv2.initUndistortRectifyMap(mtx_left,dist_left,R1, newcameramtx_left,(w,h),5)
    # mapx,mapy = cv2.initUndistortRectifyMap(mtx_left,dist_left,R1, newcameramtx_left,(w,h),5)
    dst_left = cv2.remap(imgL,mapx,mapy,cv2.INTER_LINEAR)
    
    # crop the image
    # x_l,y_l,w_l,h_l = roi_left
    # dst_left = dst_left[y_l:y_l+h_l, x_l:x_l+w_l]
    
    
    newcameramtx_right, roi_right=cv2.getOptimalNewCameraMatrix(mtx_right,dist_right,(w,h),1,(w,h))
    # dst_right = cv2.undistort(imgR, mtx_right, dist_right, R2, newcameramtx_right)
    # mapx,mapy = cv2.initUndistortRectifyMap(mtx_right,dist_right,R2,None,(w,h),5)
    mapx,mapy = cv2.initUndistortRectifyMap(mtx_right,dist_right,R2,newcameramtx_right,(w,h),5)
    dst_right = cv2.remap(imgR,mapx,mapy,cv2.INTER_LINEAR)

    # crop the image
    # x,y,w,h = roi_right
    # dst_right = dst_right[y:y+h, x:x+w]
    
    cv2.imwrite(os.path.join(SAVE_PATH, 'left3_new2.jpg'), dst_left)
    cv2.imwrite(os.path.join(SAVE_PATH, 'right3_new2.jpg'), dst_right)
    


def Clc3D(lpnt, rpnt, mtx_left, Q):
    '''
        lpnt:需要检测的左边点, 2*n numpy
        rpnt:需要检测的右边点, 2*n numpy
    '''
     #    point4D = cv2.triangulatePoints(P1, P2, lpnt, rpnt)
    d = lpnt[0] - rpnt[0]
    f = mtx_left[0, 0]
    w = Q[3, 2] * d + Q[3,3]
    
    x = -(lpnt[0] + Q[0,3])/ w
    y = -(lpnt[1] + Q[1,3]) / w
    z = -f / w
    return [x, y, z]
     # return point4D

def main(width, height, squareSize, in_path, out_path, target_path):
    '''
        width: 棋盘格的水平方向个数
        height: 棋盘格的竖直方向个数
        squareSize: 棋盘格在现实中的大小
        in_path: 标定相机内参的照片，文件夹内left为左边相机的照片，right为右边相机的照片
        out_path: 标定相机外参的照片
        target_path: 存放待检测图片的文件夹
    '''
    boardSize = (width, height)

    try:
        npzfile = np.load('calibrate_left.npz')
        mtx_left = npzfile['mtx']
        dist_left = npzfile['dist']
    except IOError:
        mtx_left, dist_left = calibrate(os.path.join(in_path, 'left'), boardSize, squareSize, True)
    try:
        npzfile = np.load('calibrate_right.npz')
        mtx_right = npzfile['mtx']
        dist_right = npzfile['dist']
    except IOError:
        mtx_right, dist_right = calibrate(os.path.join(in_path, 'right'),boardSize, squareSize,  False)
    print(mtx_left)
    print(dist_left)
    print(mtx_right)
    print(dist_right)
    try:
        npzfile = np.load('calibrate_stero.npz')
        R = npzfile['R']
        T = npzfile['T']
        R1 = npzfile['R1']
        # T1 = npzfile['T1']
        P1 = npzfile['P1']
        R2 = npzfile['R2']
        # T2 = npzfile['T2']
        P2 = npzfile['P2']
        T = npzfile['T']
        Q = npzfile['Q']
        shape = npzfile['shape']
    except IOError:

        objpoints, imgpoints, shape = WriteImagePointsInStero()

        R1, R2, P1, P2, R, T, Q = strtoCalibrate(objpoints, imgpoints, mtx_left, dist_left, mtx_right, dist_right, shape)

    correctImage(os.path.join(target_path, 'left.jpg'), os.path.join(target_path, 'right.jpg'), shape, mtx_left, dist_left, R1, mtx_right, dist_right, R2)      
    lpnt = np.array([3370,1844 ]).reshape(-1, 1)
    rpnt = np.array([1547,1616 ]).reshape(-1, 1)
    print(Clc3D(lpnt, rpnt, mtx_left, Q))
#    


if __name__ == '__main__':
    
    main(9, 6, 2.65, r'20200103Pattern2/one', r'20200103Pattern2/two', '20200103Pattern2')
    


